//! Transformations to expressions.

use std::collections::HashMap;

use stencil_tysys::EnvExpr;

use super::TaggedExpr;

/// Describes a mapping of env exprs to other env exprs.
pub struct EnvMapping {
    map: HashMap<EnvExpr, EnvExpr>,
}

impl EnvMapping {
    pub fn new(map: HashMap<EnvExpr, EnvExpr>) -> Self {
        Self { map }
    }

    fn transform_env_expr(&self, ee: EnvExpr) -> EnvExpr {
        self.map.get(&ee).cloned().unwrap_or(ee)
    }

    /// Transforms an expression to an equivalent expression but replacing env
    /// expressions with substitutions where applicable.
    pub fn apply<T>(&self, e: TaggedExpr<T>) -> TaggedExpr<T> {
        match e {
            TaggedExpr::Const(_, _) => e,
            TaggedExpr::Env(t, ee) => TaggedExpr::Env(t, self.transform_env_expr(ee)),
            TaggedExpr::Monop(t, a, b) => TaggedExpr::Monop(t, a, Box::new(self.apply(*b))),
            TaggedExpr::Binop(t, a, b, c) => {
                TaggedExpr::Binop(t, a, Box::new(self.apply(*b)), Box::new(self.apply(*c)))
            }
            TaggedExpr::Cond(t, a, b, c) => TaggedExpr::Cond(
                t,
                Box::new(self.apply(*a)),
                Box::new(self.apply(*b)),
                Box::new(self.apply(*c)),
            ),
            TaggedExpr::Cast(t, a, b) => TaggedExpr::Cast(t, a, Box::new(self.apply(*b))),
            TaggedExpr::BufferLen(t, a) => TaggedExpr::BufferLen(t, Box::new(self.apply(*a))),
            TaggedExpr::BufferRead(t, a, b, c) => TaggedExpr::BufferRead(
                t,
                Box::new(self.apply(*a)),
                Box::new(self.apply(*b)),
                Box::new(self.apply(*c)),
            ),
            TaggedExpr::BufferReadInt(t, a, b, c) => {
                TaggedExpr::BufferReadInt(t, Box::new(self.apply(*a)), Box::new(self.apply(*b)), c)
            }
        }
    }
}
